﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using Dia2Lib;
using Mono.Cecil.PE;

namespace NHook
{
	static class DIAExtensions
	{
		public static IEnumerable<IDiaTable> ToEnumerable(this IDiaEnumTables tables)
		{
			IDiaTable currentTable = null;
			uint celt = 0;
			while(true)
			{
				tables.Next(1, ref currentTable, ref celt);
				if(celt != 1)
					break;
				yield return currentTable;
			}
		}

		public static IEnumerable<IDiaSymbol> ToEnumerable(this IDiaEnumSymbols symbols)
		{
			IDiaSymbol currentSymbol = null;
			uint celt = 0;
			while(true)
			{
				symbols.Next(1, out currentSymbol, out celt);
				if(celt != 1)
					break;
				yield return currentSymbol;
			}
		}
	}

	public enum SymbolType
	{
		Export,
		Pdb
	}
	public class SymbolInfo
	{

		public string Name
		{
			get;
			set;
		}
		public RVA RVA
		{
			get;
			set;
		}
		public SymbolType Type
		{
			get;
			set;
		}
		public override string ToString()
		{
			return Enum.GetName(typeof(SymbolType), Type) + " - " + Name;
		}
	}

	public class SymbolPath
	{
		public string CachePath
		{
			get;
			set;
		}
		public string SymbolServer
		{
			get;
			set;
		}
		public override string ToString()
		{
			return "SRV*" + CachePath + "*" + SymbolServer;
		}
	}
	public class SymbolManager
	{
		#region Errors
		private readonly static uint E_PDB_OK;
		private readonly static uint E_PDB_USAGE;
		private readonly static uint E_PDB_OUT_OF_MEMORY;
		private readonly static uint E_PDB_FILE_SYSTEM;
		private readonly static uint E_PDB_NOT_FOUND;
		private readonly static uint E_PDB_INVALID_SIG;
		private readonly static uint E_PDB_INVALID_AGE;
		private readonly static uint E_PDB_PRECOMP_REQUIRED;
		private readonly static uint E_PDB_OUT_OF_TI;
		private readonly static uint E_PDB_NOT_IMPLEMENTED;
		private readonly static uint E_PDB_V1_PDB;
		private readonly static uint E_PDB_FORMAT;
		private readonly static uint E_PDB_LIMIT;
		private readonly static uint E_PDB_CORRUPT;
		private readonly static uint E_PDB_TI16;
		private readonly static uint E_PDB_ACCESS_DENIED;
		private readonly static uint E_PDB_ILLEGAL_TYPE_EDIT;
		private readonly static uint E_PDB_INVALID_EXECUTABLE;
		private readonly static uint E_PDB_DBG_NOT_FOUND;
		private readonly static uint E_PDB_NO_DEBUG_INFO;
		private readonly static uint E_PDB_INVALID_EXE_TIMESTAMP;
		private readonly static uint E_PDB_RESERVED;
		private readonly static uint E_PDB_DEBUG_INFO_NOT_IN_PDB;
		private readonly static uint E_PDB_SYMSRV_BAD_CACHE_PATH;
		private readonly static uint E_PDB_SYMSRV_CACHE_FULL;
		private readonly static uint E_PDB_MAX;
		private readonly static Dictionary<uint, string> _ErrorMessages = new Dictionary<uint, string>();
		static SymbolManager()
		{
			E_PDB_OK = (((((uint)1 << 31) | ((uint)0x6d << 16)) | (uint)1));
			E_PDB_USAGE = (E_PDB_OK + 1);
			E_PDB_OUT_OF_MEMORY = (E_PDB_USAGE + 1);
			E_PDB_FILE_SYSTEM = (E_PDB_OUT_OF_MEMORY + 1);
			E_PDB_NOT_FOUND = (E_PDB_FILE_SYSTEM + 1);
			E_PDB_INVALID_SIG = (E_PDB_NOT_FOUND + 1);
			E_PDB_INVALID_AGE = (E_PDB_INVALID_SIG + 1);
			E_PDB_PRECOMP_REQUIRED = (E_PDB_INVALID_AGE + 1);
			E_PDB_OUT_OF_TI = (E_PDB_PRECOMP_REQUIRED + 1);
			E_PDB_NOT_IMPLEMENTED = (E_PDB_OUT_OF_TI + 1);
			E_PDB_V1_PDB = (E_PDB_NOT_IMPLEMENTED + 1);
			E_PDB_FORMAT = (E_PDB_V1_PDB + 1);
			E_PDB_LIMIT = (E_PDB_FORMAT + 1);
			E_PDB_CORRUPT = (E_PDB_LIMIT + 1);
			E_PDB_TI16 = (E_PDB_CORRUPT + 1);
			E_PDB_ACCESS_DENIED = (E_PDB_TI16 + 1);
			E_PDB_ILLEGAL_TYPE_EDIT = (E_PDB_ACCESS_DENIED + 1);
			E_PDB_INVALID_EXECUTABLE = (E_PDB_ILLEGAL_TYPE_EDIT + 1);
			E_PDB_DBG_NOT_FOUND = (E_PDB_INVALID_EXECUTABLE + 1);
			E_PDB_NO_DEBUG_INFO = (E_PDB_DBG_NOT_FOUND + 1);
			E_PDB_INVALID_EXE_TIMESTAMP = (E_PDB_NO_DEBUG_INFO + 1);
			E_PDB_RESERVED = (E_PDB_INVALID_EXE_TIMESTAMP + 1);
			E_PDB_DEBUG_INFO_NOT_IN_PDB = (E_PDB_RESERVED + 1);
			E_PDB_SYMSRV_BAD_CACHE_PATH = (E_PDB_DEBUG_INFO_NOT_IN_PDB + 1);
			E_PDB_SYMSRV_CACHE_FULL = (E_PDB_SYMSRV_BAD_CACHE_PATH + 1);
			E_PDB_MAX = (E_PDB_SYMSRV_CACHE_FULL + 1);
			_ErrorMessages[E_PDB_OK] = "No error";
			_ErrorMessages[E_PDB_USAGE] = "E_PDB_USAGE";
			_ErrorMessages[E_PDB_OUT_OF_MEMORY] = "Out of memory";
			_ErrorMessages[E_PDB_FILE_SYSTEM] = "E_PDB_FILE_SYSTEM";
			_ErrorMessages[E_PDB_NOT_FOUND] = "Failed to open the file, or the file has an invalid format";
			_ErrorMessages[E_PDB_INVALID_SIG] = "Signature does not match";
			_ErrorMessages[E_PDB_INVALID_AGE] = "Age does not match";
			_ErrorMessages[E_PDB_PRECOMP_REQUIRED] = "E_PDB_PRECOMP_REQUIRED";
			_ErrorMessages[E_PDB_OUT_OF_TI] = "E_PDB_OUT_OF_TI";
			_ErrorMessages[E_PDB_NOT_IMPLEMENTED] = "E_PDB_NOT_IMPLEMENTED";
			_ErrorMessages[E_PDB_V1_PDB] = "E_PDB_V1_PDB";
			_ErrorMessages[E_PDB_FORMAT] = "Obsolete format";
			_ErrorMessages[E_PDB_LIMIT] = "E_PDB_LIMIT";
			_ErrorMessages[E_PDB_CORRUPT] = "E_PDB_CORRUPT";
			_ErrorMessages[E_PDB_TI16] = "E_PDB_TI16";
			_ErrorMessages[E_PDB_ACCESS_DENIED] = "Access denied";
			_ErrorMessages[E_PDB_ILLEGAL_TYPE_EDIT] = "E_PDB_ILLEGAL_TYPE_EDIT";
			_ErrorMessages[E_PDB_INVALID_EXECUTABLE] = "Invalid executable";
			_ErrorMessages[E_PDB_DBG_NOT_FOUND] = "Debug file not found";
			_ErrorMessages[E_PDB_NO_DEBUG_INFO] = "E_PDB_NO_DEBUG_INFO";
			_ErrorMessages[E_PDB_INVALID_EXE_TIMESTAMP] = "E_PDB_INVALID_EXE_TIMESTAMP";
			_ErrorMessages[E_PDB_RESERVED] = "E_PDB_RESERVED";
			_ErrorMessages[E_PDB_DEBUG_INFO_NOT_IN_PDB] = "E_PDB_DEBUG_INFO_NOT_IN_PDB";
			_ErrorMessages[E_PDB_SYMSRV_BAD_CACHE_PATH] = "E_PDB_SYMSRV_BAD_CACHE_PATH";
			_ErrorMessages[E_PDB_SYMSRV_CACHE_FULL] = "E_PDB_SYMSRV_CACHE_FUL";
			_ErrorMessages[E_PDB_MAX] = "E_PDB_MAX";
		}
		string GetErrorMessage(uint error)
		{
			string err;
			if(_ErrorMessages.TryGetValue(error, out err))
				return err;
			return "Unknown error " + error;
		}
		#endregion

		private ProcessDebugger _Debugger;
		Dictionary<string, DiaSource> _PDBs = new Dictionary<string, DiaSource>();
		Dictionary<string, ExportEntry[]> _Exports = new Dictionary<string, ExportEntry[]>();

		DiaSource FindSource(string module)
		{
			DiaSource source;
			if(_PDBs.TryGetValue(module, out source))
				return source;
			throw new InvalidOperationException("PDB file of module " + module + " not loaded");
		}


		public SymbolManager(ProcessDebugger processDebugger)
		{
			this._Debugger = processDebugger;
			SymbolPath = new SymbolPath()
			{
				CachePath = "SymbolsCache",
				SymbolServer = "http://msdl.microsoft.com/download/symbols"
			};
		}

		public bool ManualLoading
		{
			get;
			set;
		}

		public SymbolPath SymbolPath
		{
			get;
			set;
		}


		public SymbolInfo FromName(string moduleName, string name)
		{
			return FromName(moduleName, name, null);
		}
		public SymbolInfo FromName(string moduleName, string name, SymTagEnum? type)
		{
			var exports = FindExports(moduleName);
			if(exports != null)
			{
				var export = exports.FirstOrDefault(e => name.Equals(e.Name, StringComparison.InvariantCultureIgnoreCase));
				if(export != null)
					return new SymbolInfo()
					{
						Name = export.Name,
						RVA = export.RVA,
						Type = SymbolType.Export
					};
			}
			_Debugger.EnsureProcessLoaded();
			DiaSource diaSource = FindSource(moduleName);
			IDiaSession session;
			IDiaEnumTables tables;
			diaSource.openSession(out session);
			session.getEnumTables(out tables);
			return tables.ToEnumerable()
				.OfType<IDiaEnumSymbols>()
				.SelectMany(s => s.ToEnumerable())
				.Where(s => s.name == name && (type == null || (uint)type.Value == s.symTag))
				.Select(s => new SymbolInfo()
				{
					Name = s.name,
					RVA = new RVA(s.relativeVirtualAddress),
					Type = SymbolType.Pdb
				})
				.FirstOrDefault();
		}

		private ExportEntry[] FindExports(string moduleName)
		{
			ExportEntry[] entries;
			if(_Exports.TryGetValue(moduleName, out entries))
				return entries;
			return null;
		}

		public void LoadSymbolsFromExe(string exePath)
		{
			LoadPDB(exePath);
			LoadExports(exePath);
		}

		private void LoadExports(string exePath)
		{
			var img = Image.ReadFromFile(exePath);
			var name = Path.GetFileName(exePath);
			_Exports.Add(name, img.GetExports());
		}

		private void LoadPDB(string exePath)
		{
			if(File.Exists(exePath))
			{
				var source = LoadSymbols(exePath, false);
				if(source == null)
					return;
				_PDBs.Add(Path.GetFileName(exePath), source);
				return;
			}
		}

		private DiaSource LoadSymbols(string exePath, bool throws)
		{
			if(!File.Exists(exePath))
				throw new FileNotFoundException(exePath);
			var diaSource = COMHelper.RegisterIfNotExistAndCreate(() => new DiaSource(), "msdia110.dll");
			try
			{
				diaSource.loadDataForExe(exePath, SymbolPath == null ? null : SymbolPath.ToString(), null);
			}
			catch(COMException comEx)
			{
				if(throws)
					throw new FileNotFoundException(GetErrorMessage((uint)comEx.ErrorCode), comEx);
				return null;
			}

			return diaSource;
		}
	}
}
